Regional Ocean: Animations of Surface Fields

Regional Ocean: Animations of Surface Fields#

Hide code cell source

%load_ext autoreload
%autoreload 2
import xarray as xr
import os
from IPython.display import HTML

import regional_utils as utils

Hide code cell source

CESM_output_dir = ""  # "CROCODILE_tutorial_nwa12_MARBL"
case_name = ""  # "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output/"

## No timeseries or base case used in this notebook
ts_dir = None
base_case_output_dir = None
base_case_name = None

## As regional domains vary so much in purpose, simulation length, and extent, we don't want to assume a minimum duration
## Thus, we ignore start and end dates and simply reduce/output over the whole time frame for all of the examples given.
start_date = None  # "0001-01-01"
end_date = None  # "0101-01-01"
base_start_date = None  # "0001-01-01"
base_end_date = None  # "0101-01-01"

obs_data_dir = None

savefigs = False
fig_output_dir = None

serial = False  # use dask LocalCluster

lc_kwargs = {}
# Parameters
case_name = "CROCODILE_tutorial_nwa12_MARBL"
base_case_name = None
CESM_output_dir = (
    "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output/"
)
start_date = "2000-01-01"
end_date = "2000-11-01"
base_start_date = ""
base_end_date = ""
lc_kwargs = {"threads_per_worker": 1}
savefigs = True
fig_output_dir = None
serial = True
subset_kwargs = {}
product = "/glade/work/ajanney/CUPiD/examples/regional_ocean/computed_notebooks//ocn/Regional_Ocean_Animations.ipynb"
# Parameters
case_name = "CROCODILE_tutorial_nwa12_MARBL"
base_case_name = "CROCODILE_tutorial_nwa12_MARBL"
CESM_output_dir = "/glade/campaign/cesm/development/cross-wg/diagnostic_framework/CESM_output_for_testing/"
start_date = "2000-01-01"
end_date = "2000-11-01"
base_start_date = ""
base_end_date = ""
ts_dir = "/glade/derecho/scratch/ajanney/archive/"
lc_kwargs = {"threads_per_worker": 1}
serial = True
savefigs = True
fig_output_dir = None
subset_kwargs = {}
product = "/glade/work/ajanney/CUPiD/examples/regional_ocean/computed_notebooks//ocn/Regional_Ocean_Report_Card.ipynb"

Hide code cell source

OUTDIR = f"{CESM_output_dir}/{case_name}/ocn/hist/"
print("Output directory is:", OUTDIR)
Output directory is: /glade/campaign/cesm/development/cross-wg/diagnostic_framework/CESM_output_for_testing//CROCODILE_tutorial_nwa12_MARBL/ocn/hist/

Open Datasets and Define Paths#

Hide code cell source

case_output_dir = os.path.join(CESM_output_dir, case_name, "ocn", "hist")

# Xarray time decoding things
time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)

## Static data includes hgrid, vgrid, bathymetry, land/sea mask
static_data = xr.open_mfdataset(
    os.path.join(case_output_dir, f"*static.nc"),
    decode_timedelta=True,
    decode_times=time_coder,
)

## Surface Data
sfc_data = xr.open_mfdataset(
    os.path.join(case_output_dir, f"*sfc*.nc"),
    decode_timedelta=True,
    decode_times=time_coder,
)

## Monthly Domain Data
# monthly_data = xr.open_mfdataset(
#     os.path.join(case_output_dir, f"*z*.nc"),
#     decode_timedelta=True,
#     decode_times=time_coder,
# )

## Image/Gif Output Directory
if fig_output_dir is None:
    image_output_dir = os.path.join(
        "/glade/derecho/scratch/",
        os.environ["USER"],
        "archive",
        case_name,
        "ocn",
        "cupid_images",
    )
else:
    image_output_dir = os.path.join(fig_output_dir, case_name, "ocn", "cupid_images")
if not os.path.exists(image_output_dir):
    os.makedirs(image_output_dir)
print("Image output directory is:", image_output_dir)
Image output directory is: /glade/derecho/scratch/ajanney/archive/CROCODILE_tutorial_nwa12_MARBL/ocn/cupid_images
## Apply time boundaries
if len(start_date) > 0 and len(end_date) > 0:
    import cftime

    calendar = sfc_data.time.encoding.get("calendar", "standard")

    calendar_map = {
        "gregorian": cftime.DatetimeProlepticGregorian,
        "noleap": cftime.DatetimeNoLeap,
    }

    CFTime = calendar_map.get(calendar, cftime.DatetimeGregorian)
    y, m, d = [int(i) for i in start_date.split("-")]
    start_date_time = CFTime(y, m, d)
    y, m, d = [int(i) for i in end_date.split("-")]
    end_date_time = CFTime(y, m, d)

    sfc_data = sfc_data.sel(time=slice(start_date_time, end_date_time))

The GIF!#

variables = ["SSH", "tos"]  # "tos", "sos", "speed", "SSV", "SSU"]
dataset = sfc_data
if dataset["time"].size > 60:  # only gen gifs for 60 days at a time
    dataset = dataset.isel(time=slice(0, 60))
for gif_variable in variables:

    field = dataset[gif_variable]

    coords = utils.chooseGeoCoords(field.dims)
    areacello = utils.chooseAreacello(field.dims)

    anim = utils.create2DFieldAnimation(
        field,
        latitude=static_data[coords["latitude"]],
        longitude=static_data[coords["longitude"]],
        iter_dim="time",
        interval=150,
        save=savefigs,
        save_path=image_output_dir,
    )
from matplotlib import rcParams

rcParams["animation.embed_limit"] = (
    200 * 1024 * 1024
)  # constrains max size of HTML to be displayed in notebook

HTML(anim.to_jshtml())